# SPDX-FileCopyrightText: 2024 German Aerospace Center (DLR.e.V.)
# SPDX-FileContributor: Ferdinand Rewicki <ferdinand.rewicki@dlr.de>
#
# SPDX-License-Identifier: Apache-2.0

"""Implementation of the SAAI."""
from operator import itemgetter

import numpy as np


def saai_score(X, y_pred, k, theta=0.5, lambda_=0.5):
    """Calculate the SAAI.

    Parameters
    ----------
    X : List of Tuples        The start and end index of the anomalies
    y_pred : ArrayLike        The predicted clusters for X
    k : int                   The number of clusters
    theta : float             Minimum overlap to consider sequences as
                              synchronized in [0,1]
    lambda_ : float           The weight of the main term  in [0,1]

    Returns
    -------
    saai                The saai value
    A_S                 The number of synchronized anomalies
    A_S_star            The number of synchronized anomalies in the same cluster.
    n_pseudo_clusters   The number of pseudo clusters

    """
    A_S, A_S_star = 0, 0
    E, S = [], []

    # 1. Create Events
    for i, (start, end) in enumerate(X):
        E.append(
            {
                "sid": i,
                "type": "START",
                "time": start,
                "start": start,
                "end": end,
                "c": y_pred[i],
            }
        )
        E.append(
            {
                "sid": i,
                "type": "END",
                "time": end,
                "start": start,
                "end": end,
                "c": y_pred[i],
            }
        )

    # 2. Sort events
    E.sort(key=itemgetter("time", "type"))
    # pp.pprint(E)

    # 3. Build A_S and A_S_star
    for e in E:
        if e["type"] == "START":
            for s in S:
                v = (min(s["end"], e["end"]) - max(s["start"], e["start"])) / (
                    max(s["end"], e["end"]) - min(s["start"], e["start"])
                )
                if v > theta:
                    A_S += 1
                    if e["c"] == s["c"]:
                        A_S_star += 1
            S.append(e)
        else:
            index_to_remove = S.index(
                {
                    "sid": e["sid"],
                    "type": "START",
                    "time": e["start"],
                    "start": e["start"],
                    "end": e["end"],
                    "c": e["c"],
                }
            )
            del S[index_to_remove]

    n_pseudo_clusters = (np.unique(y_pred, return_counts=True)[1] == 1).sum()

    main_term = A_S_star / A_S if A_S_star > 0 else 0.0

    score = lambda_ * main_term + (1 - lambda_) * ((k - 1 - n_pseudo_clusters)/k)

    return score, A_S, A_S_star, n_pseudo_clusters
